{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "多标签分类\n",
    "\n",
    "可能有多个标签，可以选择"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from datasets import load_dataset, load_from_disk, concatenate_datasets\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "\n",
    "from prompt import get_industry_trans_func\n",
    "from utils import glm4_vllm, load_obj\n",
    "from setting import StaticValues\n",
    "from parse import reason_label_parse\n",
    "from collections import Counter\n",
    "from bert_train import bert_cls_trans, BertCLS, bert_tokenize_func\n",
    "\n",
    "model_name = \"bert-base-chinese\"\n",
    "\n",
    "os.environ[\"HTTP_PROXY\"] = \"http://127.0.0.1:7890\"\n",
    "os.environ[\"HTTPS_PROXY\"] = \"http://127.0.0.1:7890\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "name = \"biomedical\"\n",
    "sv = StaticValues(name=name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_multi_dataset = load_from_disk(sv.bert_config.bert_multi_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['企业名称', '经营范围', '一级行业分类', '二级行业分类', '三级行业分类', 'reason', 'label', '原料及装备', '生物医药生产', '生物医药流通', '其他'],\n",
       "        num_rows: 30000\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['企业名称', '经营范围', '一级行业分类', '二级行业分类', '三级行业分类', 'reason', 'label'],\n",
       "        num_rows: 319438\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bert_multi_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jie/anaconda3/envs/llm/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "new_dataset = bert_cls_trans(name, bert_multi_dataset[\"train\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['企业名称', '经营范围', '一级行业分类', '二级行业分类', '三级行业分类', 'reason', 'label', '原料及装备', '生物医药生产', '生物医药流通', '其他', 'industry_info', 'input_ids', 'token_type_ids', 'attention_mask'],\n",
       "    num_rows: 6916\n",
       "})"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    sv.bert_config.model_name, num_labels=len(sv.LABEL_NAME)\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(sv.bert_config.model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['企业名称', '经营范围', '一级行业分类', '二级行业分类', '三级行业分类', 'reason', 'label', '原料及装备', '生物医药生产', '生物医药流通', '其他', 'industry_info', 'input_ids', 'token_type_ids', 'attention_mask'],\n",
       "    num_rows: 6916\n",
       "})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jie/anaconda3/envs/llm/lib/python3.10/site-packages/transformers/training_args.py:1525: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "bert_cls = BertCLS(model=model, tokenizer=tokenizer, train_dataset=new_dataset,\n",
    "                   output_dir=sv.bert_config.output_multi_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_cls.train(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "719616fbaf104b1aa0534a2f457ec573",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/55 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "array([0, 0, 0, ..., 3, 3, 1])"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bert_cls.predict(new_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jie/anaconda3/envs/llm/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n",
      "/home/jie/anaconda3/envs/llm/lib/python3.10/site-packages/transformers/training_args.py:1525: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    \"/home/jie/gitee/pku_industry/general/test_output/best_model\",\n",
    "    num_labels=len(sv.LABEL_NAME),\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(sv.bert_config.model_name)\n",
    "bert_cls = BertCLS(model=model, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['企业名称', '经营范围', '一级行业分类', '二级行业分类', '三级行业分类', 'reason', 'label'],\n",
       "    num_rows: 319438\n",
       "})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bert_tokenize_func(item):\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\"bert-base-chinese\")\n",
    "    return tokenizer(\n",
    "        item[\"industry_info\"],\n",
    "        max_length=512,\n",
    "        truncation=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_dataset = bert_multi_dataset[\"test\"]\n",
    "pred_dataset = pred_dataset.map(get_industry_trans_func(\"industry_info\", \"{industry_info}\"))\n",
    "pred_dataset = pred_dataset.map(bert_tokenize_func, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_ids = bert_cls.predict(pred_dataset)\n",
    "pred_dataset = pred_dataset.add_column(\"pred_label\", pred_ids)\n",
    "pred_multi_csv = os.path.join(sv.home_folder, \"bert_multi_pred.csv\")\n",
    "pred_dataset.to_csv(pred_multi_csv, index=False)"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": ""
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 导出预测数据集"
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "cell_type": "code",
   "source": [
    "import pandas as pd\n",
    "from datasets import load_dataset, load_from_disk, concatenate_datasets\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "cell_type": "code",
   "source": [
    "name = \"biomedical\"\n",
    "sv = StaticValues(name=name)\n",
    "\n",
    "pred_multi_csv = os.path.join(sv.home_folder, \"bert_multi_pred.csv\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "cell_type": "code",
   "source": "pred_multi_csv",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "cell_type": "code",
   "source": "pd.read_csv(pred_multi_csv)",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": ""
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
